# table_from_boxes_overlay.py
from typing import List, Tuple, Any
from PIL import Image, ImageDraw, ImageFont

Box = Tuple[float, float, float, float]  # (x1, y1, x2, y2)

# ---------- 与之前相同的聚类 & 构网格函数 ---------- #
def _cluster_lines(coords: List[float], tol: int = 5) -> List[float]:
    coords_sorted = sorted(coords)
    clusters = [[coords_sorted[0]]]
    for c in coords_sorted[1:]:
        if abs(c - clusters[-1][-1]) <= tol:
            clusters[-1].append(c)
        else:
            clusters.append([c])
    return [sum(g) / len(g) for g in clusters]

def build_table_structure(cell_boxes: List[Box], tol: int = 5):
    xs, ys = [], []
    for x1, y1, x2, y2 in cell_boxes:
        xs.extend([x1, x2]); ys.extend([y1, y2])

    x_lines, y_lines = _cluster_lines(xs, tol), _cluster_lines(ys, tol)
    n_rows, n_cols = len(y_lines)-1, len(x_lines)-1
    grid: List[List[Any]] = [[None]*n_cols for _ in range(n_rows)]

    for idx, (x1, y1, x2, y2) in enumerate(cell_boxes):
        row_top = max(i for i, y in enumerate(y_lines) if y <= y1+tol)
        row_bot = min(i for i, y in enumerate(y_lines) if y >= y2-tol) - 1
        col_l   = max(j for j, x in enumerate(x_lines) if x <= x1+tol)
        col_r   = min(j for j, x in enumerate(x_lines) if x >= x2-tol) - 1

        rs, cs = row_bot-row_top+1, col_r-col_l+1
        grid[row_top][col_l] = dict(id=idx, rowspan=rs, colspan=cs)
        for r in range(row_top, row_top+rs):
            for c in range(col_l, col_l+cs):
                if (r, c) != (row_top, col_l):
                    grid[r][c] = "SPAN"
    return x_lines, y_lines, grid

def grid_to_html(grid, header_row: int = 0) -> str:
    html = ["<table>"]
    for r, row in enumerate(grid):
        html.append("  <tr>")
        c = 0
        while c < len(row):
            cell = row[c]
            if cell in (None, "SPAN"):
                c += 1; continue
            tag = "th" if r == header_row else "td"
            attrs = []
            if cell["rowspan"] > 1: attrs.append(f'rowspan="{cell["rowspan"]}"')
            if cell["colspan"] > 1: attrs.append(f'colspan="{cell["colspan"]}"')
            attr_str = " "+ " ".join(attrs) if attrs else ""
            html.append(f'    <{tag}{attr_str}>cell_{cell["id"]}</{tag}>')
            c += cell["colspan"]
        html.append("  </tr>")
    html.append("</table>")
    return "\n".join(html)

# ---------- 可视化（叠加在原图） ---------- #
def overlay_on_jpg(img_path: str,
                   cell_boxes: List[Box],
                   x_lines: List[float],
                   y_lines: List[float],
                   grid: List[List[Any]]):

    base = Image.open(img_path).convert("RGB")
    font = ImageFont.load_default()

    # ---------- ① 画 OCR box ----------
    img_boxes = base.copy()
    db = ImageDraw.Draw(img_boxes, "RGBA")
    for box in cell_boxes:
        db.rectangle(box, outline=(255, 0, 0, 150), width=2)
    img_boxes.save("./output/debug_boxes.jpg")

    # ---------- ② 画分割线 ----------
    img_lines = base.copy()
    dl = ImageDraw.Draw(img_lines)
    w, h = img_lines.size
    for y in y_lines:
        dl.line([(0, y), (w, y)], fill=(255, 0, 0), width=3)
    for x in x_lines:
        dl.line([(x, 0), (x, h)], fill=(255, 0, 0), width=3)
    img_lines.save("./output/debug_lines.jpg")

    # ---------- ③ 画真实网格 + 合并框 ----------
    img_grid = base.copy()
    dg = ImageDraw.Draw(img_grid)
    for r, row in enumerate(grid):
        for c, cell in enumerate(row):
            x1, y1 = x_lines[c], y_lines[r]
            x2, y2 = x_lines[c + 1], y_lines[r + 1]

            # 单格蓝框
            dg.rectangle([x1, y1, x2, y2], outline=(255, 0, 0), width=1)

            # 合并单元格粗框
            if cell and cell != "SPAN":
                if cell["rowspan"] > 1 or cell["colspan"] > 1:
                    rs, cs = cell["rowspan"], cell["colspan"]
                    dg.rectangle([x1, y1, x_lines[c + cs], y_lines[r + rs]],
                                 outline=(255, 0, 0), width=3)
                dg.text((x1 + 2, y1 + 2), str(cell["id"]),
                        fill=(0, 0, 0), font=font)
    img_grid.save("./output/debug_grid.jpg")

# ---------------------- DEMO ---------------------- #
if __name__ == "__main__":
    # ←←把下方两处替换成你的数据和图片路径→→
    image_path = "/home/fengjie/GITEE/pdf2md_without_ocr/inputs/3a417bf2cc6a33e21bf2cd995da6898bd1de20d6a35b24157b27a77244616b1d.jpg"      # 原始 JPG 表格
    cell_box_list = [[14.017175137996674, 0.09666550159454346, 262.3449945449829, 1410.249863743782], [262.901650428772, 0.41444939374923706, 849.5864496231079, 44.52292740345001], [849.9718866348267, 0.12429302930831909, 1025.8128900527954, 44.52399933338165], [1026.1165399551392, 0.09666550159454346, 1226.0342645645142, 44.89133942127228], [1226.1474237442017, 0.09666550159454346, 1402.1023797988892, 44.92103683948517], [1402.3847284317017, 0.09666550159454346, 1571.7211542129517, 44.59300720691681], [262.31328678131104, 44.496846318244934, 849.5461664199829, 97.7072001695633], [849.8588495254517, 44.684041142463684, 1025.8554315567017, 97.54522812366486], [1025.8507318496704, 45.04889404773712, 1226.2491083145142, 97.68416702747345], [1226.0956659317017, 45.1354113817215, 1402.2574090957642, 97.73102676868439], [1402.2665643692017, 44.70793259143829, 1572.5291376113892, 97.78488266468048], [262.27307987213135, 97.55204117298126, 849.5317010879517, 149.74908554553986], [849.9361200332642, 97.47808945178986, 1025.763451576233, 149.77145493030548], [1025.8937005996704, 97.59712326526642, 1226.1388788223267, 149.66146957874298], [1226.0396356582642, 97.63482010364532, 1402.0529413223267, 149.60513412952423], [1402.1159296035767, 97.66991531848907, 1573.0116815567017, 149.5915538072586], [262.55281925201416, 149.7386943101883, 849.3900995254517, 202.72811996936798], [849.8722162246704, 149.79609787464142, 1025.725121498108, 202.59359848499298], [1025.921166419983, 149.73660385608673, 1226.0712518692017, 202.40203964710236], [1226.1133661270142, 149.59407150745392, 1401.9881219863892, 202.43695175647736], [1402.0872430801392, 149.6677714586258, 1573.4247674942017, 202.5161143541336], [262.72640323638916, 202.53547775745392, 849.3805170059204, 254.9191142320633], [849.8540277481079, 202.43382370471954, 1025.7043085098267, 254.73872482776642], [1025.8186883926392, 202.22181808948517, 1226.1070184707642, 254.64906418323517], [1225.9944696426392, 202.26693832874298, 1401.9772577285767, 254.62837326526642], [1402.0122919082642, 202.35543930530548, 1573.4345331192017, 254.61766159534454], [262.94812870025635, 254.7468272447586, 849.3955926895142, 307.7077189683914], [849.7719354629517, 254.65321457386017, 1025.7554559707642, 307.53147995471954], [1025.894066810608, 254.50254929065704, 1226.3224725723267, 307.3851786851883], [1226.1281366348267, 254.5170146226883, 1402.0308465957642, 307.3031474351883], [1402.0762567520142, 254.5526133775711, 1572.5445184707642, 307.4695292711258], [263.1728448867798, 307.57301437854767, 849.5251703262329, 359.50965988636017], [849.8159418106079, 307.4016581773758, 1025.8772821426392, 359.6342021226883], [1025.968773841858, 307.2664042711258, 1226.4733514785767, 359.63414108753204], [1226.2044305801392, 307.1675273180008, 1402.1818475723267, 359.69944870471954], [1402.1360712051392, 307.36900436878204, 1571.8460321426392, 359.7589579820633], [263.23799991607666, 359.40977585315704, 849.5995111465454, 412.2357951402664], [849.8364496231079, 359.5476847887039, 1025.9164056777954, 412.43043625354767], [1026.0168085098267, 359.53553879261017, 1226.5259637832642, 412.6202861070633], [1226.2354364395142, 359.4880839586258, 1402.3139276504517, 412.67576706409454], [1402.2663202285767, 359.71934616565704, 1571.8660516738892, 412.7759562730789], [263.094277381897, 412.0077372789383, 849.6819696426392, 464.3314982652664], [849.8496942520142, 412.21522629261017, 1025.947777748108, 464.57554733753204], [1026.0580072402954, 412.4156962633133, 1226.4812860488892, 464.81590378284454], [1226.3338251113892, 412.4998332262039, 1402.3699579238892, 464.81874191761017], [1402.3705682754517, 412.54185593128204, 1572.0288934707642, 464.6796122789383], [263.0907373428345, 464.37132370471954, 849.6074457168579, 517.416367650032], [849.8163080215454, 464.62364304065704, 1025.8455438613892, 517.7001200914383], [1025.9417352676392, 464.8238993883133, 1226.3725214004517, 517.884141087532], [1226.2379999160767, 464.82103073596954, 1402.3926630020142, 517.7797099351883], [1402.3960809707642, 464.7190715074539, 1572.0170526504517, 517.7272197008133], [262.7734155654907, 517.4394999742508, 849.5493402481079, 569.7790385484695], [849.8540277481079, 517.6602030992508, 1025.8046503067017, 569.9747782945633], [1025.9119501113892, 517.7265483140945, 1226.3250360488892, 570.0790873765945], [1226.3251581192017, 517.7308818101883, 1402.3025751113892, 569.892686009407], [1402.3860712051392, 517.672959446907, 1572.1003046035767, 569.6822978258133], [262.7854242324829, 569.5785990953445, 849.5581293106079, 622.607285618782], [849.8421869277954, 569.7723857164383, 1025.8487787246704, 622.6717997789383], [1026.035668373108, 569.7537089586258, 1226.2221307754517, 622.6962138414383], [1226.2664422988892, 569.5823222398758, 1402.1042108535767, 622.6662455797195], [1402.2953729629517, 569.5657206773758, 1572.1850214004517, 622.705186009407], [262.62664127349854, 622.3049174547195, 849.6877679824829, 674.8132182359695], [849.8282709121704, 622.5632792711258, 1025.8938837051392, 675.062852025032], [1025.997582435608, 622.4946147203445, 1226.1458368301392, 675.075547337532], [1226.1703729629517, 622.4740458726883, 1402.0589227676392, 674.8832865953445], [1402.2587518692017, 622.4828959703445, 1572.3542108535767, 674.8070536851883], [262.4701623916626, 674.7854472398758, 849.7268915176392, 728.0863505601883], [849.8196649551392, 675.0849467515945, 1025.8988885879517, 728.0026713609695], [1025.971459388733, 675.0932475328445, 1226.1869745254517, 727.9471904039383], [1226.2034540176392, 674.9609843492508, 1402.2417840957642, 727.7964335680008], [1402.3933954238892, 674.8396464586258, 1572.3891229629517, 727.7344218492508], [262.6897211074829, 727.870225071907, 849.8275995254517, 780.1854716539383], [849.7331171035767, 727.8889018297195, 1025.9320917129517, 780.1658183336258], [1025.9024286270142, 727.8566142320633, 1226.2278680801392, 780.2299662828445], [1226.2211542129517, 727.762559056282, 1402.3123407363892, 779.9785014390945], [1402.5053339004517, 727.590928196907, 1572.4695672988892, 779.825059056282], [262.8910150527954, 779.982041478157, 849.7721185684204, 833.1269389390945], [849.8097772598267, 780.0940409898758, 1025.8902826309204, 832.9127665758133], [1025.8644647598267, 780.1013652086258, 1226.2502069473267, 833.0421611070633], [1226.0866327285767, 779.9164897203445, 1402.2875604629517, 832.931992650032], [1402.4030389785767, 779.725693821907, 1572.1136102676392, 832.9324198961258], [262.69666385650635, 832.8535014390945, 849.7396478652954, 884.9323588609695], [849.7232294082642, 832.797227025032, 1025.9624872207642, 884.970566868782], [1025.820336341858, 832.917588353157, 1226.2027215957642, 885.1164408922195], [1226.0125360488892, 832.8519755601883, 1402.1964960098267, 885.1008158922195], [1402.2498407363892, 832.6743632555008, 1572.1034784317017, 884.9858256578445], [262.8012628555298, 884.7507182359695, 849.7025384902954, 938.120469212532], [849.7237787246704, 884.9267436265945, 1025.8485345840454, 937.9340068101883], [1025.843285560608, 885.058457493782, 1226.1731805801392, 938.0224467515945], [1226.0277948379517, 885.0820780992508, 1402.2776727676392, 938.0150614976883], [1402.2104120254517, 884.9579325914383, 1571.8463983535767, 938.0844584703445], [262.68059635162354, 938.0002909898758, 849.6169672012329, 990.0034037828445], [849.7868280410767, 937.8246928453445, 1025.688988685608, 990.0367900133133], [1025.8265008926392, 937.8791362047195, 1226.2353143692017, 990.0327616930008], [1226.0586786270142, 937.914170384407, 1402.3140497207642, 989.9792338609695], [1402.3406610488892, 937.8861552476883, 1571.8835077285767, 989.9836283922195], [262.693398475647, 989.9687358140945, 849.4790887832642, 1042.793564915657], [849.7674188613892, 990.023301243782, 1025.667748451233, 1042.8910990953445], [1025.8846063613892, 990.0113993883133, 1226.2731561660767, 1042.743271946907], [1226.1754999160767, 990.024766087532, 1402.2019891738892, 1042.765488743782], [1402.2416620254517, 989.9301615953445, 1571.9283075332642, 1043.005967259407], [262.7277307510376, 1042.647568821907, 849.4651117324829, 1094.821396946907], [849.7580194473267, 1042.616562962532, 1025.7120599746704, 1094.969834446907], [1025.9892206192017, 1042.5866557359695, 1226.2088251113892, 1094.965195775032], [1226.0236444473267, 1042.5783549547195, 1402.1923456192017, 1094.960801243782], [1402.0785760879517, 1042.7753764390945, 1571.7426385879517, 1094.9667826890945], [262.7885217666626, 1094.7079936265945, 849.4409418106079, 1147.760605931282], [849.7857294082642, 1094.8740092515945, 1025.7350702285767, 1147.737900853157], [1025.959496498108, 1094.9191752672195, 1226.2528924942017, 1147.7221537828445], [1226.0397577285767, 1094.869980931282, 1402.1498651504517, 1147.6948100328445], [1402.0999383926392, 1094.861680150032, 1571.6623163223267, 1147.816758275032], [262.88619327545166, 1147.5954447984695, 849.5767450332642, 1199.7382670640945], [849.8205804824829, 1147.6139994859695, 1025.8300409317017, 1199.883652806282], [1025.9532098770142, 1147.601426243782, 1226.3575067520142, 1199.909287571907], [1226.1294794082642, 1147.5268412828445, 1402.2626581192017, 1199.962998509407], [1402.1756219863892, 1147.6816264390945, 1571.5642938613892, 1199.995469212532], [262.8688135147095, 1199.723008275032, 849.5423822402954, 1252.583115696907], [849.8546380996704, 1199.837998509407, 1025.8404169082642, 1252.687363743782], [1025.9007806777954, 1199.7643901109695, 1226.4549188613892, 1252.8835307359695], [1226.2152948379517, 1199.823350071907, 1402.2946405410767, 1253.030869603157], [1402.2454462051392, 1199.9250346422195, 1571.9540643692017, 1252.994248509407], [262.8475275039673, 1252.402451634407, 849.5625848770142, 1304.7280131578445], [849.8741693496704, 1252.5099955797195, 1025.8429193496704, 1304.857041478157], [1025.9364862442017, 1252.7226420640945, 1226.4762811660767, 1305.092881321907], [1226.2035760879517, 1252.8342143297195, 1402.3457880020142, 1305.080186009407], [1402.3793573379517, 1252.8515483140945, 1571.9569940567017, 1304.904160618782], [262.5334711074829, 1304.696152806282, 849.5987787246704, 1357.8771830797195], [849.8015985488892, 1304.886338353157, 1025.730248451233, 1357.921738743782], [1025.8509759902954, 1305.0768901109695, 1226.3531122207642, 1358.1525737047195], [1226.2190790176392, 1305.128037571907, 1402.1969842910767, 1357.9477397203445], [1402.2489862442017, 1304.9767924547195, 1572.0880975723267, 1358.0632182359695], [262.527871131897, 1357.824082493782, 849.4834222793579, 1409.9391947984695], [849.8338861465454, 1357.8993998765945, 1025.7071161270142, 1410.1745463609695], [1025.8536615371704, 1358.009873509407, 1226.2660760879517, 1410.2953959703445], [1226.1980829238892, 1357.9191752672195, 1402.1265497207642, 1410.118027806282], [1402.2819452285767, 1357.995957493782, 1572.0889520645142, 1409.9575053453445], [14.625112473964691, 1409.684678196907, 262.53369998931885, 1779.092881321907], [262.3168878555298, 1409.5698100328445, 849.6063470840454, 1462.938584446907], [849.8522577285767, 1409.965928196907, 1025.7390985488892, 1462.816269993782], [1025.8998041152954, 1409.972275853157, 1226.1430292129517, 1462.9391947984695], [1226.1059198379517, 1409.885117650032, 1402.1105585098267, 1462.8925639390945], [1402.4171991348267, 1409.852646946907, 1572.5279169082642, 1463.019150853157], [262.30738162994385, 1462.9111186265945, 849.7311639785767, 1515.043809056282], [849.7996454238892, 1462.732773900032, 1025.8008661270142, 1515.131455540657], [1025.8830194473267, 1462.7504740953445, 1225.9848260879517, 1515.2851420640945], [1226.0198602676392, 1462.750107884407, 1402.1602411270142, 1515.013291478157], [1402.3725214004517, 1462.8332377672195, 1572.4176874160767, 1515.024277806282], [262.16980838775635, 1515.1039897203445, 849.7064447402954, 1568.161973118782], [849.8095941543579, 1515.2997905015945, 1025.806176185608, 1568.040146946907], [1025.895531654358, 1515.3151713609695, 1226.0232782363892, 1568.0912944078445], [1226.0063104629517, 1515.1681987047195, 1402.2045526504517, 1567.8823100328445], [1402.3924188613892, 1515.001816868782, 1571.7947626113892, 1567.7814799547195], [262.089150428772, 1568.056992650032, 849.6639032363892, 1620.099961400032], [849.7697992324829, 1567.8554545640945, 1025.812035560608, 1620.186875462532], [1025.821557044983, 1567.951279759407, 1226.0928583145142, 1620.2692729234695], [1226.0667352676392, 1567.782578587532, 1402.2860956192017, 1620.093857884407], [1402.4338006973267, 1567.6828471422195, 1571.6635370254517, 1620.005723118782], [262.3068017959595, 1620.035264134407, 849.6400384902954, 1673.101182103157], [849.7974481582642, 1620.120713353157, 1025.7720575332642, 1672.891465306282], [1025.8128900527954, 1620.134873509407, 1226.1801385879517, 1673.174180150032], [1226.0501337051392, 1619.9745951890945, 1402.1864862442017, 1673.2045756578445], [1402.2851190567017, 1619.838486790657, 1571.7688837051392, 1673.2631694078445], [262.237832069397, 1672.938340306282, 849.7248773574829, 1725.165391087532], [849.7094964981079, 1672.7309428453445, 1025.7083368301392, 1725.186143040657], [1025.8324823379517, 1672.9570170640945, 1226.0965204238892, 1725.3381205797195], [1225.9187860488892, 1673.095811009407, 1402.1103143692017, 1725.2302104234695], [1402.1346063613892, 1673.0456401109695, 1571.8647089004517, 1725.360215306282], [262.44727420806885, 1725.0483256578445, 849.5475702285767, 1778.5463725328445], [849.7307977676392, 1725.052842259407, 1025.7279291152954, 1778.1918803453445], [1025.7641229629517, 1725.160264134407, 1226.1054315567017, 1778.192490696907], [1225.8571405410767, 1725.1562358140945, 1402.0522089004517, 1778.1083842515945], [1402.1150751113892, 1725.196885228157, 1571.8949823379517, 1778.1882182359695], [14.882627129554749, 1778.720078587532, 262.4616632461548, 2036.2216655015945], [262.43428897857666, 1778.279160618782, 849.3297357559204, 1830.1794291734695], [849.6551141738892, 1778.0297709703445, 1025.6177606582642, 1830.1628276109695], [1025.732689857483, 1778.0598002672195, 1226.0058221817017, 1830.2272807359695], [1225.8260126113892, 1777.980576634407, 1402.0602655410767, 1830.2101908922195], [1402.0601434707642, 1777.9841166734695, 1571.7369012832642, 1830.157822728157], [262.4190454483032, 1829.9509135484695, 849.2694330215454, 1883.089463353157], [849.6375970840454, 1830.1169291734695, 1025.5565423965454, 1882.9697123765945], [1025.7770013809204, 1830.1950541734695, 1225.9740839004517, 1882.930039525032], [1225.9180536270142, 1830.1984721422195, 1401.9661493301392, 1882.9450541734695], [1402.1136102676392, 1830.116074681282, 1571.5481805801392, 1883.146836400032], [262.51966190338135, 1882.856309056282, 849.4061517715454, 1935.0080424547195], [849.5936517715454, 1882.809678196907, 1025.5174798965454, 1935.180283665657], [1025.701683998108, 1882.710312962532, 1226.0603876113892, 1935.110703587532], [1225.9361200332642, 1882.7888041734695, 1401.9300165176392, 1935.0956889390945], [1402.0170526504517, 1882.9418803453445, 1571.5585565567017, 1935.1179057359695], [262.5926752090454, 1934.7529155015945, 849.3782587051392, 1987.9250346422195], [849.5325555801392, 1935.0549174547195, 1025.5041131973267, 1987.7612162828445], [1025.6094598770142, 1935.0309916734695, 1226.1924676895142, 1987.782578587532], [1225.9455194473267, 1934.975205540657, 1401.9273309707642, 1987.825303196907], [1401.8949823379517, 1934.960312962532, 1571.4150018692017, 1988.303330540657], [262.678017616272, 1987.5185405015945, 848.9947137832642, 2029.449814915657], [849.2901239395142, 1987.8251811265945, 1025.445824623108, 2030.1818705797195], [1025.5132684707642, 1987.6357280015945, 1226.2553339004517, 2031.1757670640945], [1225.8439569473267, 1987.724961400032, 1401.9746942520142, 2031.3312846422195], [1401.5594110488892, 1988.226426243782, 1571.3975458145142, 2030.502549290657]]

    x_lines, y_lines, grid = build_table_structure(cell_box_list, tol=5)
    overlay_on_jpg(image_path, cell_box_list, x_lines, y_lines, grid)

    html = grid_to_html(grid, header_row=0)
    with open("./output/table.html", "w", encoding="utf-8") as f:
        f.write(html)
    print("⎋ 调试图已输出：debug_boxes.jpg / debug_lines.jpg / debug_grid.jpg")
    print("⎋ HTML 保存为 table.html")